program program_psr
use madx_ptc_module
use pointer_lattice
implicit none

    interface
       subroutine build_tiny_package_example(PSR,MIS)
         use madx_ptc_module
         use pointer_lattice
         implicit none
         type(layout), target :: PSR
         logical(lp) mis
       end subroutine build_tiny_package_example
    end interface

    interface
       subroutine build_PSR(PSR,MIS)
         use madx_ptc_module
         use pointer_lattice
         implicit none
         type(layout), target :: PSR
         logical(lp) mis
       end subroutine build_PSR
    end interface

type(layout), pointer:: psr
type(damap) id,m,a_t,a_cs
type(real_8) ray(6)
real(dp) closed_orbit(6),DX_AVERAGE_DJ,beta,phase_x
type(internal_state) state
type(fibre), pointer :: f
type(normalform) normal
integer :: no1=0,np1=0 ,ind(2)=0,i 
type(taylor) z_i,average_floquet,average_xp,PHASE_ADVANCE(2)
logical(lp) mis,toy

call ptc_ini

psr=>m_u%start

Write(6,*) " small misalignments and rotations in lattice ? input t or f "
read(5,*) mis
Write(6,*) " For toy Lattice  input t or PSR input f "
read(5,*) toy

write(6,*) " Give  i for <x_i> "
read(5,*) ind(1) 

if(toy) then
 call build_tiny_package_example(PSR,mis)
else 
 call build_PSR(PSR,mis)
endif

!!!  Phase space of Taylor maps will be (x,px,y,py,delta)
state=delta0


do while(no1<=0) 

   write(6,'(a90)') " with order=1, this code will produce only the dispersion dclosed_orbit/dx_5         "
   write(6,*) " "
   write(6,'(a90)') " in general you get  dclosed_orbit^n/(dx_5^n1 dJ_1^n2 dJ_2^n3) (n1+2*n2+2*n3 = order)"
   write(6,*) " "
   if(toy.and. ind(1) ==1) then
    write(6,'(a90)') " input an order  ---->  order > 1  if you want comparison with analytical for <x>   "
    write(6,*) " "
    read(5,*) no1
    
     if(no1>1) then
        write(6,*) " Since you selected an order greater than 1,"
        write(6,*) " You may want to study the code z_my_nonlinear_twiss_average_x.f90 "
        write(6,*) " in build_my_own_tiny_package\my_main_programs "
        write(6,*) " "
    endif  
  else
   write(6,*) " input an order > 0 "
   read(5,*) no1
  endif
enddo

np1=0    !!!  No system parameters in the map

call init(state,no1,np1) ! PTC and FPP are properly initialized

call alloc(id,m,A_T,a_cs)  
call alloc(ray)
call alloc(normal)
call alloc(z_i,average_floquet,average_xp)
call alloc(PHASE_ADVANCE)

closed_orbit=0.d0 ! initial guess for closed orbit

call find_orbit(psr,closed_orbit,1,state,1.d-5)

write(6,'(a16,4(1x,g12.5))') " closed orbit = ",closed_orbit(1:4)

id=1;   ray=closed_orbit+id;   ! ray= closed orbit + identity map

call track(psr,ray,1,state) ! One turn map is computed via the ray

!!! The ray is really the result of a truncated power series algebra
!!! The map we are interested to compute is around the closed orbit: it is really part of a diffential algebra
m=ray   !   The ray is "officially" turned into a "damap"  (DA in Berz's honour)

normal = m  ! The map is normalised   



z_i=id%v(ind(1))


!!! This averages an arbitrary function over the orbit
!!! Internally this routine suffers in clarity from the absence of a Complex TPSA  
!!! Input = function z_i and transformation normal%A_T   -> Normal form = normal%A_T^-1 o Map o normal%A_T
call AVERAGE(z_i,normal%A_T,average_floquet,average_xp,use_J=.true.)


!!! Output = average_floquet -> average as a function of J_i= (x_i^2 + p_i^2)/2
!!! J_1 is in position 1 and J_2 in position 3
Write(6,*) " average_floquet from output of AVERAGE"
call print(average_floquet,6)

!!! Output = average_xp -> average as a function of x,px,y,py in the ORIGINAL VARIABLES

Write(6,*) " average_xp from output of AVERAGE"
call print(average_xp,6)

!!!! Let us do a test: construct average_xp using average_floquet and normal%A_T^-1 
!!!! First create a special map

m=0

m%v(1)= (id%v(1)**2+id%v(2)**2)/2.d0   !   J_1
m%v(3)= (id%v(3)**2+id%v(4)**2)/2.d0   !   J_2

m=m*normal%A_T**(-1)   ! going in the original variables I_1= J_1 *normal%A_T**(-1)    SEE SECTION ON INVARIANTS

average_xp=0

average_xp=average_floquet*m

Write(6,*) " average_xp reconstructed using output of AVERAGE"
call print(average_xp,6)

if(toy.and.no1>0.and.ind(1)==1) then
 write(6,*) " "
 Write(6,*) " Since you have selected the toy lattice and , <x>," 
 Write(6,*) " We will compare with an analytical calculation for d<x>/dJ_x"
 Write(6,*) "  without errors the agreement with TPSA/FPP should be perfect "
 Write(6,*) "  Doing a Twiss to integrate the formula for d<x>/dJ_x "
 Write(6,*) " read the chapter on this topic if you are lost "
 write(6,*) " "
 
 !  Here starts a universal Twiss loop
 f=>psr%start               
 DX_AVERAGE_DJ=0.d0
 a_t=normal%a_t
  call CANONIZE( a_t,A_cs,PHASE_ADVANCE)  ! We canonize the the transformation A_T into A_CS
 PHASE_ADVANCE(1)=0.0_dp

 ray=A_cs+closed_orbit   ! The ray is initialize as the map A_cs+closed_orbit

 phase_x=0.0_dp

 do i=1,psr%n
  call track(psr,ray,i,i+1,state)

 !!! a very dirty but correct implementation of the beta function calculation, valid with coupling
  a_t=ray
  beta=(A_T%v(1).sub.'10')*(A_T%v(1).sub.'10')+(A_T%v(1).sub.'01')*(A_T%v(1).sub.'01')   ! beta function extracted the dirty way
 !!!
 
  call CANONIZE( a_t,A_cs,PHASE_ADVANCE)    ! A_t= A_cs o R(phase_advance_i_i+1)    Universal definition of phase advance

 phase_x=PHASE_ADVANCE(1) 

  closed_orbit=ray
  ray=A_cs+closed_orbit    ! Ray is rebooted into a canonical form to insure that PHASE_ADVANCE contains only the incremental phase advance

 !!! d<x_normal>/dJ is summed (from analytical perturbation theory) 
 if(f%mag%p%nmul>=3) then
   DX_AVERAGE_DJ=BETA**1.5_DP*f%mag%BN(3)/2.0_DP *(-SIN(phase_x*TWOPI)+SIN((phase_x-NORMAL%TUNE(1))*TWOPI)) &
     /(1.0_DP-COS(NORMAL%TUNE(1)*TWOPI)) + DX_AVERAGE_DJ
 endif
 f=>f%next
 enddo
 DX_AVERAGE_DJ=DX_AVERAGE_DJ*SQRT(BETA)  ! d<x>/dJ= d<x_normal>/dJ sqrt(beta)
 if(mis) then
   write(6,*) "The agreement will not be perfect because of misalignments "
 else
  write(6,*) "The agreement should be perfect "
 endif
 write(6,*) " Analytical result ", DX_AVERAGE_DJ
  DX_AVERAGE_DJ=average_floquet.sub.'1'
 write(6,*) " FPP result        ", DX_AVERAGE_DJ

endif



write(6,*) "   "
write(6,*) " hit return to terminate program " 
write(6,*) "   "
pause 
call kill(id,m,A_T,a_cs)  
call kill(ray)
call kill(normal)
call kill(z_i,average_floquet,average_xp)
call kill(PHASE_ADVANCE)

end program program_psr


!=================================================================

subroutine  build_tiny_package_example(PSR,mis)
use madx_ptc_module
use pointer_lattice
implicit none

type(layout), target :: PSR

real(dp) :: ang, brho, kd, kf, Larc,SG(6),cut,H
type(fibre) ::  sf, L, qd, qf
type(layout) :: cell
logical(lp) mis
!-----------------------------------

call make_states(.false.)
call update_states

madlength = .false.
exact_model = .false.
madkind2 = drift_kick_drift

ang = (twopi / 20.d0)
Larc = 1.d0
brho = 1.2d0 * (Larc / ang)
h=  twopi/20/Larc
call set_mad(brho = brho, method = 2, step = 100)

kf = 1.d0
kd = -1.d0

L=DRIFT("L",Larc)
qf = sbend("QF", Larc, ang).q.kf  ! <-- PUTS QUADRUPOLE KF
qd = sbend("QD", Larc, ang).q.kD  ! <-- PUTS QUADRUPOLE KD

sf = sextupole("SF", 0.d0, 1.d0) 


cell = QF+SF+L+QD+SF+L

PSR = 10 * cell
PSR = .ring.PSR

call survey(PSR)

if(mis) then
 sig=1.d-5
 cut=4.d0
 call MESS_UP_ALIGNMENT(PSR,SIG,cut)
endif
end subroutine build_tiny_package_example

subroutine  build_PSR(PSR,mis)
use madx_ptc_module
use pointer_lattice
implicit none

type(layout), target :: PSR

real(dp) :: ang, brho, kd, kf, Larc,SG(6),cut
type(fibre) :: b, d1, d2, qd, qf
type(layout) :: cell
logical(lp) mis
!-----------------------------------

call make_states(.false.)
exact_model = .true.
default = default + nocavity  
call update_states
madlength = .false.

ang = (twopi * 36.d0 / 360.d0)
Larc = 2.54948d0
brho = 1.2d0 * (Larc / ang)
call set_mad(brho = brho, method = 2, step = 10)
madkind2 = drift_kick_drift

kf =  2.72d0 / brho
kd = -1.92d0 / brho

d1 = drift("D1", 2.28646d0)
d2 = drift("D2", 0.45d0)
qf = quadrupole("QF", 0.5d0, kf)
qd = quadrupole("QD", 0.5d0, kd)
b  = rbend("B", Larc, ang)
cell = d1 + qd + d2 + b + d2 + qf + d1

PSR = 10 * cell
PSR = .ring.PSR

call survey(PSR)

if(mis) then
 sig=1.d-5
 cut=4.d0
 call MESS_UP_ALIGNMENT(PSR,SIG,cut)
endif
end subroutine build_PSR
